#!/usr/bin/env python3
"""
环境设置脚本

在PyCharm中运行此脚本来设置项目环境
"""

import os
import sys
import subprocess
from pathlib import Path


def check_python_version():
    """检查Python版本"""
    if sys.version_info < (3, 8):
        print("错误: 需要Python 3.8或更高版本")
        return False
    print(f"Python版本: {sys.version}")
    return True


def check_dependencies():
    """检查依赖包"""
    required_packages = [
        'torch', 'torchvision', 'opencv-python', 'numpy', 
        'matplotlib', 'scikit-image', 'PyYAML', 'tqdm'
    ]
    
    missing_packages = []
    
    for package in required_packages:
        try:
            __import__(package.replace('-', '_'))
            print(f"✓ {package}")
        except ImportError:
            missing_packages.append(package)
            print(f"✗ {package}")
    
    if missing_packages:
        print(f"\n缺少以下包: {missing_packages}")
        print("请运行: pip install -r requirements.txt")
        return False
    
    return True


def check_cuda():
    """检查CUDA可用性"""
    try:
        import torch
        if torch.cuda.is_available():
            print(f"✓ CUDA可用: {torch.cuda.device_count()}个GPU")
            for i in range(torch.cuda.device_count()):
                print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
            return True
        else:
            print("✗ CUDA不可用，将使用CPU")
            return False
    except ImportError:
        print("✗ PyTorch未安装")
        return False


def create_directories():
    """创建必要的目录"""
    directories = [
        'data/raw/echocardiogram',
        'data/processed/segmentation',
        'data/annotations/manual',
        'data/annotations/semi_auto',
        'data/preprocessing',
        'data/augmentation',
        'checkpoints/sam',
        'checkpoints/prompt',
        'checkpoints/memory',
        'checkpoints/periodic',
        'checkpoints/quality_aware',
        'logs/training',
        'logs/inference',
        'logs/evaluation',
        'experiments/baseline',
        'experiments/periodic_prompt',
        'experiments/quality_aware',
        'experiments/memory_mask',
        'outputs/predictions',
        'outputs/visualization'
    ]
    
    for directory in directories:
        Path(directory).mkdir(parents=True, exist_ok=True)
        print(f"✓ 创建目录: {directory}")


def setup_python_path():
    """设置Python路径"""
    project_root = Path(__file__).parent.parent
    src_path = project_root / 'src'
    
    if str(src_path) not in sys.path:
        sys.path.insert(0, str(src_path))
        print(f"✓ 添加Python路径: {src_path}")
    
    # 设置环境变量
    os.environ['PYTHONPATH'] = str(src_path)
    print(f"✓ 设置PYTHONPATH: {src_path}")


def test_imports():
    """测试关键模块导入"""
    try:
        from src.utils.logger import get_logger
        from src.utils.config import load_config
        from src.utils.metrics import calculate_all_metrics
        from src.utils.visualization import plot_training_curves
        print("✓ 工具模块导入成功")
        return True
    except ImportError as e:
        print(f"✗ 模块导入失败: {e}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("SAM-Echo-Video 环境设置")
    print("=" * 60)
    
    # 检查Python版本
    print("\n1. 检查Python版本...")
    if not check_python_version():
        return False
    
    # 检查依赖包
    print("\n2. 检查依赖包...")
    if not check_dependencies():
        print("\n请先安装依赖包:")
        print("pip install -r requirements.txt")
        return False
    
    # 检查CUDA
    print("\n3. 检查CUDA...")
    cuda_available = check_cuda()
    
    # 创建目录
    print("\n4. 创建项目目录...")
    create_directories()
    
    # 设置Python路径
    print("\n5. 设置Python路径...")
    setup_python_path()
    
    # 测试导入
    print("\n6. 测试模块导入...")
    if not test_imports():
        print("模块导入测试失败，请检查项目结构")
        return False
    
    # 环境信息
    print("\n" + "=" * 60)
    print("环境设置完成!")
    print("=" * 60)
    print(f"项目根目录: {Path(__file__).parent.parent}")
    print(f"Python版本: {sys.version}")
    print(f"CUDA可用: {cuda_available}")
    print(f"Python路径: {sys.path[0]}")
    
    print("\n下一步:")
    print("1. 在PyCharm中配置运行配置")
    print("2. 运行数据预处理脚本")
    print("3. 开始模型训练")
    
    return True


if __name__ == "__main__":
    success = main()
    if not success:
        sys.exit(1)
